/////////////////////////////////////////////////////////////////////////////
// Original code from libhdfs3. Copyright (c) 2013 - 2014, Pivotal Inc.
// All rights reserved. Author: Zhanwei Wang
/////////////////////////////////////////////////////////////////////////////
//  Modifications by Kumo Inc.
// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//



#include <ktest/ktest.h>

#include <kmhdfs/hdfs.h>
#include <kmhdfs/client/kerberos_name.h>
#include <kmhdfs/common/exception.h>
#include <kmhdfs/common/exception_internal.h>
#include <tests/mock/test_util.h>

#include <krb5/krb5.h>

#ifndef TEST_HDFS_PREFIX
#define TEST_HDFS_PREFIX "./"
#endif

#define BASE_DIR TEST_HDFS_PREFIX"/testSecureHACInterface/"

using namespace Hdfs;
using namespace Hdfs::Internal;

static std::string ExtractPrincipalFromTicketCache(
    const std::string & cachePath) {
    krb5_context cxt = NULL;
    krb5_ccache ccache = NULL;
    krb5_principal principal = NULL;
    krb5_error_code ec = 0;
    std::string errmsg, retval;
    char * priName = NULL;

    if (!cachePath.empty()) {
        if (0 != setenv("KRB5CCNAME", cachePath.c_str(), 1)) {
            THROW(HdfsIOException, "Cannot set env parameter \"KRB5CCNAME\"");
        }
    }

    do {
        if (0 != (ec = krb5_init_context(&cxt))) {
            break;
        }

        if (0 != (ec = krb5_cc_default(cxt, &ccache))) {
            break;
        }

        if (0 != (ec = krb5_cc_get_principal(cxt, ccache, &principal))) {
            break;
        }

        if (0 != (ec = krb5_unparse_name(cxt, principal, &priName))) {
            break;
        }
    } while (0);

    if (!ec) {
        retval = priName;
    } else {
        if (cxt) {
            errmsg = krb5_get_error_message(cxt, ec);
        } else {
            errmsg = "Cannot initialize kerberos context";
        }
    }

    if (priName != NULL) {
        krb5_free_unparsed_name(cxt, priName);
    }

    if (principal != NULL) {
        krb5_free_principal(cxt, principal);
    }

    if (ccache != NULL) {
        krb5_cc_close(cxt, ccache);
    }

    if (cxt != NULL) {
        krb5_free_context(cxt);
    }

    if (!errmsg.empty()) {
        THROW(HdfsIOException,
              "FileSystem: Filed to extract principal from ticket cache: %s",
              errmsg.c_str());
    }

    return retval;
}

class TestKerberosConnectC: public ::testing::Test {
public:
    TestKerberosConnectC() {
        setenv("LIBHDFS3_CONF", "function-secure.xml", 1);
        struct hdfsBuilder * bld = hdfsNewBuilder();
        assert(bld != NULL);
        hdfsBuilderSetNameNode(bld, "default");
        hdfsBuilderSetForceNewInstance(bld);
        std::stringstream ss;
        ss.imbue(std::locale::classic());
        ss << "/tmp/krb5cc_";
        ss << getuid();
        const char * userCCpath = GetEnv("LIBHDFS3_TEST_USER_CCPATH", ss.str().c_str());
        hdfsBuilderSetKerbTicketCachePath(bld, userCCpath);
        fs = hdfsBuilderConnect(bld);
        hdfsFreeBuilder(bld);

        if (fs == NULL) {
            throw std::runtime_error("cannot connect hdfs");
        }

        if (0 != hdfsCreateDirectory(fs, BASE_DIR)) {
            throw std::runtime_error("cannot create test directory");
        }

        name = KerberosName(ExtractPrincipalFromTicketCache(userCCpath));
    }

    ~TestKerberosConnectC() {
        hdfsDelete(fs, BASE_DIR, true);
        hdfsDisconnect(fs);
    }

protected:
    hdfsFS fs;
    KerberosName name;
};

TEST_F(TestKerberosConnectC, GetDelegationToken_Failure) {
    ASSERT_TRUE(NULL == hdfsGetDelegationToken(NULL, NULL));
    ASSERT_EQ(EINVAL, errno);
    ASSERT_TRUE(NULL == hdfsGetDelegationToken(fs, NULL));
    ASSERT_EQ(EINVAL, errno);
    ASSERT_TRUE(NULL == hdfsGetDelegationToken(fs, ""));
    ASSERT_EQ(EINVAL, errno);
}

TEST_F(TestKerberosConnectC, DelegationToken) {
    char * token = NULL;
    hdfsFreeDelegationToken(NULL);
    ASSERT_EQ(-1, hdfsRenewDelegationToken(NULL, NULL));
    ASSERT_EQ(EINVAL, errno);
    ASSERT_EQ(-1, hdfsRenewDelegationToken(fs, NULL));
    ASSERT_EQ(EINVAL, errno);
    ASSERT_TRUE(NULL != (token = hdfsGetDelegationToken(fs, "Unknown")));
    ASSERT_EQ(-1, hdfsRenewDelegationToken(fs, token));
    ASSERT_EQ(EACCES, errno);
    hdfsCancelDelegationToken(fs, token);
    hdfsFreeDelegationToken(token);
    ASSERT_TRUE(NULL != (token = hdfsGetDelegationToken(fs, name.getName().c_str())));
    ASSERT_GT(hdfsRenewDelegationToken(fs, token), 0);
    ASSERT_EQ(0, hdfsCancelDelegationToken(fs, token));
    ASSERT_EQ(-1, hdfsRenewDelegationToken(fs, token));
    ASSERT_EQ(EPERM, errno);
    ASSERT_EQ(-1, hdfsCancelDelegationToken(fs, token));
    ASSERT_EQ(EPERM, errno);
    hdfsFreeDelegationToken(token);
}
